from typing import List, Tuple, Union
import numpy as np
import cv2


class ImageTransform:
    """图像变换工具类"""

    @staticmethod
    def get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray:
        """获取第三个点的坐标"""
        direct = a - b
        return b + np.array([-direct[1], direct[0]], dtype=np.float32)

    @staticmethod
    def get_dir(src_point: List[float], rot_rad: float) -> List[float]:
        """获取旋转后的方向向量"""
        sn, cs = np.sin(rot_rad), np.cos(rot_rad)
        src_result = [0, 0]
        src_result[0] = src_point[0] * cs - src_point[1] * sn
        src_result[1] = src_point[0] * sn + src_point[1] * cs
        return src_result

    @staticmethod
    def distance(x1: float, y1: float, x2: float, y2: float) -> float:
        """计算两点之间的欧氏距离"""
        return np.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2)

    @classmethod
    def get_affine_transform(cls, center: np.ndarray, scale: Union[float, np.ndarray],
                             rot: float, output_size: Tuple[int, int],
                             shift: np.ndarray = np.array([0, 0], dtype=np.float32),
                             inv: bool = False) -> np.ndarray:
        """获取仿射变换矩阵"""
        if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
            scale = np.array([scale, scale], dtype=np.float32)

        scale_tmp = scale
        src_w = scale_tmp[0]
        dst_w, dst_h = output_size

        rot_rad = np.pi * rot / 180
        src_dir = cls.get_dir([0, src_w * -0.5], rot_rad)
        dst_dir = np.array([0, dst_w * -0.5], np.float32)

        src = np.zeros((3, 2), dtype=np.float32)
        dst = np.zeros((3, 2), dtype=np.float32)
        src[0, :] = center + scale_tmp * shift
        src[1, :] = center + src_dir + scale_tmp * shift
        dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
        dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir

        src[2:, :] = cls.get_3rd_point(src[0, :], src[1, :])
        dst[2:, :] = cls.get_3rd_point(dst[0, :], dst[1, :])

        if inv:
            trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
        else:
            trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))

        return trans

    @staticmethod
    def crop_quad(img: np.ndarray, position: np.ndarray) -> np.ndarray:
        """根据四边形位置裁剪图像"""
        x0, y0 = position[0][0], position[0][1]
        x1, y1 = position[1][0], position[1][1]
        x2, y2 = position[2][0], position[2][1]
        x3, y3 = position[3][0], position[3][1]

        img_width = ImageTransform.distance((x0 + x3) / 2, (y0 + y3) / 2,
                                            (x1 + x2) / 2, (y1 + y2) / 2)
        img_height = ImageTransform.distance((x0 + x1) / 2, (y0 + y1) / 2,
                                             (x2 + x3) / 2, (y2 + y3) / 2)

        corners_trans = np.zeros((4, 2), np.float32)
        corners_trans[0] = [0, 0]
        corners_trans[1] = [img_width, 0]
        corners_trans[2] = [img_width, img_height]
        corners_trans[3] = [0, img_height]

        transform = cv2.getPerspectiveTransform(position, corners_trans)
        dst = cv2.warpPerspective(img, transform, (int(img_width), int(img_height)))
        return dst
